import os
import paddle
import utils
import layers
import time
import numpy as np

#参数设置
scale = 3
imgPath = "test/val"

#预测
def predict(model,img,to_file="result.jpg"):
    imgs = paddle.to_tensor([img],dtype="float32")
    result = model(imgs)
    result = utils.channel_first_to_last(result.numpy()[0])*255
    utils.save_rgb(result,to_file)    
    return result  

imgPath = ["test/val/%s"%(i) for i in os.listdir('test/val')]
imgs = [utils.read_image_channel_first(i) for i in imgPath]
img_tensor = [i/255 for i in imgs]
img_tensor = paddle.to_tensor(img_tensor)
#channel first to last
imgs = [utils.channel_first_to_last(i) for i in imgs]
h,w,_ = imgs[0].shape
for scale in range(2,5):
    pnsr = 0.0
    ssim = 0.0
    print("resize scale %d"%(scale))
    #加载模型
    m = "%s/%s"%("model/x%d"%(scale),os.listdir("model/x%d"%(scale))[-1])
    edsr = layers.EDSR(scale)
    layer_state_dict = paddle.load(m)
    edsr.set_state_dict(layer_state_dict)
    #预测图片
    result = edsr(img_tensor)
    result = result.numpy()*255
    #保存图片
    for i in range(len(result)):
        img = utils.resize_bicubic(imgs[i],h*scale,w*scale)
        rr = utils.channel_first_to_last(result[i])
        pnsr += utils.psnr(img,rr)
        ssim += utils.ssim(img,rr)
        
        combine = np.concatenate((img,rr),1)
        utils.save_rgb(combine,"test/x%d/%d.jpg"%(scale,i))
    pnsr = pnsr/len(result)
    ssim = ssim/len(result)
    print("scale:")
    print("pnsr=%f, ssim=%f"%(pnsr,ssim))
        
  